#include "demo_cpp_tf/simple_kalman_filter.hpp"

namespace demo_cpp_tf 
{
    SimpleKalmanFilter::SimpleKalmanFilter(int state_dim, int measure_dim) 
    : n_(state_dim), m_(measure_dim),
        x_(state_dim), P_(state_dim, state_dim),
        F_(Eigen::MatrixXd::Identity(state_dim, state_dim)),
        Q_(Eigen::MatrixXd::Identity(state_dim, state_dim)),
        H_(Eigen::MatrixXd::Zero(measure_dim, state_dim)),
        R_(Eigen::MatrixXd::Identity(measure_dim, measure_dim)),
        is_nonlinear_(false) // 新增非线性标志
    {
        H_.block(0, 0, measure_dim, measure_dim).setIdentity();
    }

    void SimpleKalmanFilter::init(const Eigen::VectorXd& x0) 
    {
        x_ = x0;
        P_ = Eigen::MatrixXd::Identity(n_, n_) * 0.1;
    }

    void SimpleKalmanFilter::predict(double dt, const Eigen::VectorXd& u) 
    {
        if(is_nonlinear_) 
        {
            // 非线性预测（匀速圆周运动模型）
            // 状态变量：[x, y, z, vx, vy, vz, ω]
            Eigen::VectorXd x_pred(n_);
            
            // 获取当前状态
            double omega = x_(6);
            double theta = omega * dt;
            
            // 旋转矩阵
            Eigen::Matrix3d Rz;
            Rz << cos(theta), -sin(theta), 0,
                 sin(theta),  cos(theta), 0,
                 0,          0,          1;
            
            // 位置预测
            Eigen::Vector3d pos = x_.head<3>() + x_.segment<3>(3) * dt;
            pos = Rz * pos; // 应用旋转
            
            // 速度预测（考虑向心力）
            Eigen::Vector3d vel = x_.segment<3>(3) + Eigen::Vector3d(
                -omega * x_(4) * dt,  // ax = -ω*vy
                 omega * x_(3) * dt,  // ay = ω*vx
                 0
            );
            
            x_pred << pos, vel, omega;
            
            // 计算雅可比矩阵（F_jacobian）
            Eigen::MatrixXd F_j(n_, n_);
            F_j.setIdentity();
            
            // 位置关于速度的偏导
            F_j.block<3,3>(0,3) = Eigen::Matrix3d::Identity() * dt;
            
            // 位置关于ω的偏导
            F_j(0,6) = -dt * (x_(3)*sin(theta) + x_(4)*cos(theta));
            F_j(1,6) =  dt * (x_(3)*cos(theta) - x_(4)*sin(theta));
            
            // 速度关于ω的偏导
            F_j(3,6) = -x_(4) * dt;
            F_j(4,6) =  x_(3) * dt;
            
            // 更新协方差
            P_ = F_j * P_ * F_j.transpose() + Q_;
            x_ = x_pred;
        } 
        else 
        {
            // 原始线性预测
            x_ = F_ * x_;
            P_ = F_ * P_ * F_.transpose() + Q_;
        }
    }

    // 增强矩阵配置接口
    void SimpleKalmanFilter::setTransitionMatrix(const Eigen::MatrixXd& F) 
    {
        if(F.rows() != n_ || F.cols() != n_)
            throw std::invalid_argument("Transition matrix dimension mismatch");
        F_ = F;
    }

    void SimpleKalmanFilter::setProcessNoiseCov(const Eigen::MatrixXd& Q) 
    {
        if(Q.rows() != n_ || Q.cols() != n_)
            throw std::invalid_argument("Process noise covariance dimension mismatch");
        Q_ = Q;
    }

    // 添加非线性更新接口
    void SimpleKalmanFilter::update(const Eigen::VectorXd& z, const Eigen::MatrixXd& H) 
    {
        Eigen::MatrixXd S = H * P_ * H.transpose() + R_;
        K_ = P_ * H.transpose() * S.inverse();
        
        x_ = x_ + K_ * (z - H * x_);
        Eigen::MatrixXd I = Eigen::MatrixXd::Identity(n_, n_);
        P_ = (I - K_ * H) * P_;
    }

    // 添加非线性模式切换
    void SimpleKalmanFilter::enableNonlinearMode(bool enable) 
    {
        is_nonlinear_ = enable;
        if(enable) 
        {
            // 初始化非线性模型参数
            Q_ = Eigen::MatrixXd(n_, n_);
            Q_ << 0.1, 0,   0,   0,    0,    0,    0,
                  0,   0.1, 0,   0,    0,    0,    0,
                  0,   0,   0.1, 0,    0,    0,    0,
                  0,   0,   0,   0.5,  0,    0,    0,
                  0,   0,   0,   0,    0.5,  0,    0,
                  0,   0,   0,   0,    0,    0.1,  0,
                  0,   0,   0,   0,    0,    0,    0.2;
        }
    }

} // namespace demo_cpp_tf